1866H - Happy Sets - CodeForces Solution

#include<bits/stdc++.h>
#define int long long
using namespace std;

const int mod=998244353;
inline int ksm(int x,int b)
{
	int res=1;
	while(b)
	{
		if(b&1)	res=res*x%mod;
		x=x*x%mod;
		b>>=1;
	}
	return res;
}

int fac[300005],invFac[300005],n,k;

inline int C(int x,int y)
{
	if (y>x) return 0;
	if (x==y||y==0) return 1;
	return fac[x]*invFac[x-y]%mod*invFac[y]%mod;
}

inline void init(int x)
{
	fac[0]=fac[1]=1;
	for (int i=2;i<=x;i++) fac[i]=fac[i-1]*i%mod;
	invFac[x]=ksm(fac[x],mod-2);
	for (int i=x-1;i>=0;i--) invFac[i]=invFac[i+1]*(i+1)%mod;
}

signed main()
{
    init(300000);
	cin>>n>>k;
	int ans=0;
	for (int i=1,lst=1;i<=min(n,k);i++) 
	{
		int tmp=fac[i+1];
		tmp=tmp*ksm(i+1,k-i)%mod;
		tmp=(tmp-lst+mod)%mod;
		lst=(lst+tmp)%mod;
		tmp=tmp*C(n,i)%mod;
		tmp=tmp*fac[i]%mod;
		ans=(ans+tmp)%mod;
	}
	ans=(ans+1)%mod; 
	cout<<ans;
}